import pefile
import re
from capstone import Cs, CS_ARCH_X86, CS_MODE_32
from graphviz import Digraph
import openai_chat
import basic_info
def extract_iat(pe_file_path):
    """从 PE 文件中提取 IAT 函数名称"""
    pe = pefile.PE(pe_file_path)
    iat = {}

    # 遍历导入表
    for entry in pe.DIRECTORY_ENTRY_IMPORT:
        for imp in entry.imports:
            if imp.address:
                # 将地址转换为十六进制字符串，确保格式统一
                iat[f'0x{imp.address:08x}'] = imp.name.decode('utf-8')

    return iat

def extract_text_segment(pe_file_path):
    """提取 PE 文件的文本段"""
    pe = pefile.PE(pe_file_path)
    text_segment = None

    for section in pe.sections:
        if b'.text' in section.Name:
            text_segment = section
            break

    if text_segment:
        return text_segment.get_data()
    else:
        raise Exception("No text segment found")

def disassemble_code(code):
    """使用 Capstone 反汇编代码"""
    md = Cs(CS_ARCH_X86, CS_MODE_32)
    instructions = []
    for instruction in md.disasm(code, 0x1000):  # 0x1000 is a placeholder for the base address
        instructions.append(f'反汇编地址: 0x{instruction.address:x} 参数: {instruction.mnemonic} {instruction.op_str}')
    return instructions

def read_disassembly(file_path, iat):
    """读取反汇编结果并替换函数名"""
    instructions = []
    with open(file_path, 'r') as f:
        for line in f:
            match = re.search(r'反汇编地址: (0x[0-9a-fA-F]+) 参数: (.+)', line)
            if match:
                addr = match.group(1)
                opcode = match.group(2)
                if 'call dword ptr' in opcode:
                    # 提取地址
                    addr_match = re.search(r'\[(0x[0-9a-fA-F]+)\]', opcode)
                    if addr_match:
                        call_addr = addr_match.group(1)
                        call_addr_formatted = f'0x{int(call_addr, 16):08x}'
                        # 查找 IAT 中的函数名称
                        function_name = iat.get(call_addr_formatted, f'Unknown({call_addr})')
                        # 替换为函数名
                        opcode = opcode.replace(f'dword ptr [{call_addr}]', function_name)
                instructions.append((addr, opcode))
    return instructions

def create_control_flow_graph(instructions, iat, output_file):
    """生成控制流图并保存为文件"""
    dot = Digraph(comment='Control Flow Graph')
    dot.attr(size='10,10')

    # 添加节点
    for addr, opcode in instructions:
        dot.node(addr, f'{addr}\n{opcode}')

    # 添加边
    for i in range(len(instructions) - 1):
        current_addr, current_opcode = instructions[i]
        next_addr, next_opcode = instructions[i + 1]

        if 'jmp' in current_opcode or 'je' in current_opcode or 'jne' in current_opcode:
            jump_match = re.search(r'(0x[0-9a-fA-F]+)', current_opcode)
            if jump_match:
                jump_addr = jump_match.group(1)
                dot.edge(current_addr, jump_addr)

        dot.edge(current_addr, next_addr)

    dot.render(output_file, format='svg', cleanup=True)

if __name__ == "__main__":
    pe_file_path = r"D:\code\ai\malware-ai\malware-ai\test_exe\trage_sample\1"  # 替换为你的可执行文件路径
    output_file = "control_flow_graph"  # 输出文件名

    # 提取 IAT
    iat = extract_iat(pe_file_path)
    print("IAT:", iat)  # 打印 IAT 字典

    # 提取文本段并反汇编
    code = extract_text_segment(pe_file_path)
    instructions_list = disassemble_code(code)

    # 将反汇编结果写入文件
    instructions_file = "res.txt"
    with open(instructions_file, 'w') as f:
        for instruction in instructions_list:
            f.write(instruction + '\n')

    # 读取反汇编结果
    instructions = read_disassembly(instructions_file, iat)

    # 创建控制流图
    create_control_flow_graph(instructions, iat, output_file)
    striat = str(iat)

    basic_info.gen_trage(pe_file_path)

    res = openai_chat.chat(striat)
    f = open("res.md", "w")
    f.write(res)
    f.close()
    f = open("trage.html", "r")
    htmlcontent = f.read()
    f.close()
    if htmlcontent != "":
        trage_res = openai_chat.trage_chat()
        f = open("res_with_trage.md", "w")
        f.write(trage_res)
        f.close()